import os
import pandas as pd
from sklearn.preprocessing import LabelEncoder, StandardScaler

import numpy as np
import pandas as pd


def feature_extra(df: pd.DataFrame) -> pd.DataFrame:
    df = df[['Age',
                  'BusinessTravel',
                  'DistanceFromHome',
                  'EducationField',
                  'EnvironmentSatisfaction',
                  'Gender',
                  'JobInvolvement',
                  'JobLevel',
                  'JobRole',
                  'JobSatisfaction',
                  'MaritalStatus',
                  'MonthlyIncome',
                  'NumCompaniesWorked',
                  'OverTime',
                  'RelationshipSatisfaction',
                  'StockOptionLevel',
                  'TotalWorkingYears',
                  'TrainingTimesLastYear',
                  'WorkLifeBalance',
                  'YearsAtCompany',
                  'YearsInCurrentRole',
                  'YearsSinceLastPromotion']]

    return df


def feature_processing(df,is_train=True,scaler=None,encoder_dict = None):
    df = df.copy()
    # 数值型特征与类别型特征
    value_col = df.select_dtypes(include=['int64', 'float64']).columns
    category_col = df.select_dtypes(include=['object']).columns

    if is_train is True:
        # 数值特征标准化
        scaler = StandardScaler()
        df[value_col] = scaler.fit_transform(df[value_col])
        # 类别特征编码

        encoder_dict={}
        for col in category_col:
            encoder = LabelEncoder()
            df[col] = encoder.fit_transform(df[col])
            encoder_dict[col] = encoder

    else:
        # 数值特征标准化
        df[value_col] = scaler.transform(df[value_col])
        # 类别特征编码
        for col in category_col:
            encoder = encoder_dict[col]
            exist_classes = encoder.classes_.tolist()
            df[col] = df[col].apply(lambda x: encoder.transform([x])[0] if x in exist_classes else -1)

    return df, scaler, encoder_dict



